import asyncio
import logging
import sys

from pydantic import Field

from nat.builder.builder import Builder
from nat.builder.framework_enum import LLMFrameworkEnum
from nat.builder.function_info import FunctionInfo
from nat.cli.register_workflow import register_function
from nat.data_models.component_ref import FunctionRef
from nat.data_models.component_ref import LLMRef
from nat.data_models.function import FunctionBaseConfig

# 添加 fastmcp 相关导入
from fastmcp import Client
from fastmcp.client.transports import StdioTransport

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# 添加 Tavily MCP Client 配置和函数
class TavilyMCPClientConfig(FunctionBaseConfig, name="tavily_mcp_client"):
    """
    配置用于Tavily MCP客户端。
    """
    description: str = Field(description="客户端描述")
    transport_type: str = Field(default="stdio", description="传输类型: http 或 stdio")
    max_results: int = Field(default=5, description="最大搜索结果数")


@register_function(config_type=TavilyMCPClientConfig)
async def tavily_mcp_client_function(config: TavilyMCPClientConfig, builder: Builder):
    """
    实现Tavily MCP客户端功能，支持HTTP和stdio两种传输方式。
    """
    async def _tavily_mcp_client(query: str) -> str:
        """
        通过MCP客户端调用Tavily搜索服务。

        Args:
            query: 搜索查询

        Returns:
            搜索结果
        """
        if config.transport_type == "http":
            # 使用 HTTP 传输
            client = Client("http://127.0.0.1:8083/mcp")
        else:
            # 使用 stdio 传输
            transport = StdioTransport(
                command="python",
                args=["-m", "tavily_mcp_server.server"],
            )
            client = Client(transport)
        logger.info(f"Tavily MCP client: {client}")

        async with client:
            try:
                logger.info(f"Tavily MCP client 查询: {query}")
                search_docs = await client.call_tool("search", {
                    "query": query,
                    "max_results": config.max_results
                })
                logger.info(f"Tavily MCP client 查询成功！")

                # tavily_mcp_client 返回的结果是 fastmcp.client.client.CallToolResult, tavily_internet_search 返回的是拼接好的字符串；
                # 这里我们参考 tavily_internet_search 将 search_docs 简单地拼接成字符串；
                search_results = "\n\n---\n\n".join(
                    [f'<Document href="{doc["url"]}"/>\n{doc["title"]}\n{doc["content"]}\n</Document>' for doc in search_docs.structured_content['data']['results']])
                logger.info(f"Tavily MCP client 查询结果: \n{search_results}")

                return search_results
            except Exception as e:
                logger.error(f"Tavily MCP client 调用失败: {e}")
                return f"搜索失败: {str(e)}"

    try:
        yield FunctionInfo.create(
            single_fn=_tavily_mcp_client,
            description=config.description
        )
    except GeneratorExit:
        logger.warning("Tavily MCP client exited early!")
    finally:
        logger.info("Cleaning up Tavily MCP client.")


class SeedSelectionAgentConfig(FunctionBaseConfig, name="seed_selection_agent"):
    """
    配置用于种子选择的代理。
    """
    llm_name: LLMRef = Field(description="要使用的LLM名称")
    description: str = Field(description="代理的描述")
    tavily_search_tool: FunctionRef = Field(description="Tavily搜索工具引用")


@register_function(config_type=SeedSelectionAgentConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
async def seed_selection_agent_function(config: SeedSelectionAgentConfig, builder: Builder):
    """
    实现种子选择代理功能。
    """
    # 获取Tavily搜索工具
    tavily_search_fn = builder.get_function(config.tavily_search_tool)
        
    async def _seed_selection_agent(location: str, season: str) -> str:
        """
        根据地理位置和季节推荐适宜种植的农作物。

        Args:
            location: 用户所在地区
            season: 当前季节

        Returns:
            推荐的农作物品种及理由
        """
        # 构造搜索查询
        query = f"{location} {season} 适宜种植的农作物品种"
        logger.info(f"搜索查询: {query}")
        
        # 调用 Tavily 搜索工具
        # acall_invoke() 是异步调用，所以，不需要加 await;
        search_results = await tavily_search_fn.acall_invoke(query)

        # 处理搜索结果并生成推荐
        if search_results:
            # 这里可以添加更复杂的逻辑来处理搜索结果
            # 现在我们只是简单地返回搜索结果
            output_message = f"根据{location}的{season}季节条件，推荐以下农作物品种：\n\n{search_results}"
        else:
            output_message = f"抱歉，根据{location}的{season}季节条件，我没有找到相关的信息。建议您可以咨询附近的农业种植户或者专家，他们可能会提供更准确的建议。"
            
        return output_message

    try:
        yield FunctionInfo.create(
            single_fn=_seed_selection_agent, 
            description=config.description
        )
    except GeneratorExit:
        logger.warning("Seed selection agent exited early!")
    finally:
        logger.info("Cleaning up seed selection agent.")


class PlantingAgentConfig(FunctionBaseConfig, name="planting_agent"):
    """
    配置用于种植管理的代理。
    """
    llm_name: LLMRef = Field(description="要使用的LLM名称")
    description: str = Field(description="代理的描述")
    tavily_search_tool: FunctionRef = Field(description="Tavily搜索工具引用")


@register_function(config_type=PlantingAgentConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
async def planting_agent_function(config: PlantingAgentConfig, builder: Builder):
    """
    实现种植管理代理功能。
    """
    # 获取Tavily搜索工具
    tavily_search_fn = builder.get_function(config.tavily_search_tool)
    
    async def _planting_agent(crop: str, issue: str) -> str:
        """
        提供作物种植问题的解决方案。

        Args:
            crop: 作物名称
            issue: 遇到的问题

        Returns:
            问题解决方案
        """
        # 构造搜索查询
        query = f"{crop} {issue} 解决方案 种植技术" 
        logger.info(f"搜索查询: {query}")
        
        # 调用Tavily搜索工具
        search_results = await tavily_search_fn.acall_invoke(query)

        # 处理搜索结果并生成解决方案
        if search_results:
            # 这里可以添加更复杂的逻辑来处理搜索结果
            # 现在我们只是简单地返回搜索结果
            output_message = f"针对{crop}的{issue}问题，建议采取以下措施：\n\n{search_results}"
        else:
            output_message = f"抱歉，针对{crop}的{issue}问题，我没有找到相关的信息。建议您可以咨询附近的农业种植户或者专家，他们可能会提供更准确的建议。"
            
        return output_message

    try:
        yield FunctionInfo.create(
            single_fn=_planting_agent,
            description=config.description
        )
    except GeneratorExit:
        logger.warning("Planting agent exited early!")
    finally:
        logger.info("Cleaning up planting agent.")


class SalesAgentConfig(FunctionBaseConfig, name="sales_agent"):
    """
    配置用于销售预测的代理。
    """
    llm_name: LLMRef = Field(description="要使用的LLM名称")
    description: str = Field(description="代理的描述")
    tavily_search_tool: FunctionRef = Field(description="Tavily搜索工具引用")


@register_function(config_type=SalesAgentConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
async def sales_agent_function(config: SalesAgentConfig, builder: Builder):
    """
    实现销售预测代理功能。
    """
    # 获取Tavily搜索工具
    tavily_search_fn = builder.get_function(config.tavily_search_tool)
    
    async def _sales_agent(crop: str) -> str:
        """
        预测作物市场行情并制定销售策略。

        Args:
            crop: 作物名称

        Returns:
            市场预测和销售建议
        """
        # 构造搜索查询
        query = f"{crop} 市场价格趋势 销售策略"
        logger.info(f"搜索查询: {query}")
        
        # 调用Tavily搜索工具
        search_results = await tavily_search_fn.acall_invoke(query)
        
        # 处理搜索结果并生成市场预测
        if search_results:
            # 这里可以添加更复杂的逻辑来处理搜索结果
            # 现在我们只是简单地返回搜索结果
            output_message = f"关于{crop}的市场行情和销售策略：\n\n{search_results}"
        else:
            output_message = f"抱歉，关于{crop}的市场行情和销售策略，我没有找到相关的信息。建议您可以咨询附近的农业种植户或者专家，他们可能会提供更准确的建议。"
            
        return output_message

    try:
        yield FunctionInfo.create(
            single_fn=_sales_agent,
            description=config.description
        )
    except GeneratorExit:
        logger.warning("Sales agent exited early!")
    finally:
        logger.info("Cleaning up sales agent.")